#include <dirent.h>
#include <sys/stat.h>
#include <iostream>
#include <fstream>
#include <string>
#include <cstring>
#include <vector>
#include <thread>
#include <mutex>
#include <cmath>
#include <queue>

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"

#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"

#include <windows.h>


// 使用 Windows API 创建目录
bool createDirectory(const std::string &path) {
    return CreateDirectoryA(path.c_str(), NULL) || GetLastError() == ERROR_ALREADY_EXISTS;
}

// RGB转HSV，输出v（亮度）范围0-1
float rgbToV(unsigned char r, unsigned char g, unsigned char b) {
    float fr = r / 255.0f;
    float fg = g / 255.0f;
    float fb = b / 255.0f;
    float maxc = std::max(fr, std::max(fg, fb));
    //float minc = std::min(fr, std::min(fg, fb));
    float v = maxc; // 亮度V就是最大值
    return v;
}

void processImage(const std::string &inputPath, const std::string &outputPath) {
    // 加载图像
    int width, height, channels;
    unsigned char* img = stbi_load(inputPath.c_str(), &width, &height, &channels, 0);
    if (!img) {
        std::cerr << "Failed to load image: " << inputPath << std::endl;
        return;
    }

    // 配置裁剪参数
    const int target_width = 512;
    const int target_height = 512;
    const int x_start = (width - target_width) / 2;
    const int y_start = 16;

    // 创建输出缓冲区和白色模板
    std::vector<unsigned char> output(target_width * target_height * channels, 0xFF);

    // 计算有效垂直范围
    const int valid_y_start = std::max(0, y_start);
    const int valid_y_end = std::min(height, y_start + target_height);
    const int copy_height = valid_y_end - valid_y_start;

    if (copy_height > 0) {
        // 计算有效水平范围
        const int valid_x_start = std::max(0, x_start);
        const int valid_x_end = std::min(width, x_start + target_width);
        const int copy_width = valid_x_end - valid_x_start;

        // 预计算内存偏移量
        const int dst_y_offset = valid_y_start - y_start;
        const int dst_x_offset = valid_x_start - x_start;

        // 执行块拷贝
        for (int y = 0; y < copy_height; ++y) {
            const unsigned char* src_row = img + ((valid_y_start + y) * width + valid_x_start) * channels;
            unsigned char* dst_row = output.data() + ((dst_y_offset + y) * target_width + dst_x_offset) * channels;
            std::memcpy(dst_row, src_row, copy_width * channels);
        }
    }

    // 释放原始图像
    stbi_image_free(img);

    // 保存结果
    std::string outputFilePath = outputPath + "/" + inputPath.substr(inputPath.find_last_of("/\\") + 1);
    if (!stbi_write_png(outputFilePath.c_str(),
                       target_width, target_height, channels, output.data(), target_width * channels)) {
        std::cerr << "Failed to save image: " << outputFilePath << std::endl;
    }
}

// 遍历目录，获取所有 PNG 文件路径
void traverseDirectory(const std::string &inputDir, std::vector<std::string> &imagePaths) {
    DIR *dir = opendir(inputDir.c_str());
    if (!dir) {
        std::cerr << "Failed to open directory: " << inputDir << std::endl;
        return;
    }

    struct dirent *entry;
    while ((entry = readdir(dir)) != nullptr) {
        std::string name = entry->d_name;

        // 忽略当前目录 "." 和父目录 ".."
        if (name == "." || name == "..") {
            continue;
        }

        std::string fullPath = inputDir + "/" + name;

        // 检查文件类型
        struct stat statbuf;
        if (stat(fullPath.c_str(), &statbuf) == 0) {
            if (S_ISREG(statbuf.st_mode)) { // 如果是普通文件
                if (name.size() >= 4 && name.substr(name.size() - 4) == ".png") {
                    imagePaths.push_back(fullPath);
                }
            } else if (S_ISDIR(statbuf.st_mode)) { // 如果是目录
                std::cout << "Skipping directory: " << fullPath << std::endl;
            }
        } else {
            std::cerr << "Failed to stat file: " << fullPath << std::endl;
        }
    }

    closedir(dir);
}

// 每个线程的任务
void processImagesBatch(const std::vector<std::string> &imagePaths, const std::string &outputPath) {
    int icount = 0;
    for (const auto &imagePath : imagePaths) {
        processImage(imagePath, outputPath);
        icount += 1;
        if (icount % 1000 == 0)
        {
            //exit(0);
            std::cout << std::this_thread::get_id()<< " processed:" << icount << std::endl;
        }
    }
}

int main(int argc, char *argv[]) {
    if (argc != 3) {
        std::cerr << "Usage: " << argv[0] << " <input_path> <output_path>" << std::endl;
        return -1;
    }

    std::string inputPath = argv[1];
    std::string outputPath = argv[2];

    // 检查输出目录是否存在，如果不存在则创建
    struct stat statbuf;
    if (stat(outputPath.c_str(), &statbuf) != 0) {
        if (!createDirectory(outputPath)) {
            std::cerr << "Failed to create output directory: " << outputPath << std::endl;
            return -1;
        }
    }

    // 获取所有 PNG 文件路径
    std::vector<std::string> imagePaths;
    traverseDirectory(inputPath, imagePaths);

    if (imagePaths.empty()) {
        std::cerr << "No PNG images found in input directory: " << inputPath << std::endl;
        return -1;
    }

    std::cout << "find image count:" << imagePaths.size() << std::endl;

    // 启动多线程处理
    const int threadCount = 16;
    std::vector<std::thread> threads;
    std::vector<std::vector<std::string>> threadImageLists(threadCount);

    // 平均分配任务到各线程
    for (size_t i = 0; i < imagePaths.size(); ++i) {
        threadImageLists[i % threadCount].push_back(imagePaths[i]);
    }

    // 启动线程
    for (int i = 0; i < threadCount; ++i) {
        threads.emplace_back(processImagesBatch, threadImageLists[i], outputPath);
    }

    // 等待所有线程完成
    for (auto &t : threads) {
        if (t.joinable()) {
            t.join();
        }
    }

    std::cout << "Processing completed." << std::endl;
    return 0;
}
